import argparse
import pickle
import random
import time
from multiprocessing.connection import Listener
from threading import Thread



import pathlib
from configs import const
from core.embeddings import VecModel


class RPCHandler:
    def __init__(self):
        self._functions = {}

    def register_function(self, func):
        self._functions[func.__name__] = func

    def handle_connection(self, connection):
        try:
            while True:
                # Receive a message
                func_name, args= pickle.loads(connection.recv())
                # Run the RPC and send a response
                try:
                    print(func_name, args)
                    r = self._functions[func_name](args)
                    print(f'{r=}')
                    connection.send(pickle.dumps(r))
                except Exception as e:
                    connection.send(pickle.dumps(e))
        except EOFError:
            import traceback
            print(traceback.format_exc())


def rpc_server(hdlr, address, authkey):
    sock = Listener(address, authkey=authkey)
    while True:
        try:
            client = sock.accept()
            t = Thread(target=hdlr.handle_connection, args=(client,))
            t.daemon = True
            t.start()
        except Exception as e:
            print("【EXCEPTION】:", str(e))


emb_model = VecModel('text',model_path=pathlib.Path(const.MODEL_PATH_TEXT).absolute())
def emb(texts):
    try:
        
        return emb_model.embed_documents(texts)
    except Exception as e:
        return str(e)


def client(texts):
    from multiprocessing.connection import Client
    address = ('127.0.0.1', 7860)
    c=Client(address, authkey=b'infiniflow-token4kevinhu')
    c.send(pickle.dumps(('emb',texts)))
    return c.recv()

if __name__ == "__main__":
    handler = RPCHandler()
    handler.register_function(emb)

    # Run the server
    rpc_server(handler, ('0.0.0.0', 7860),
               authkey=b'infiniflow-token4kevinhu')